#include "ProjectGaussians.h"

//******************************************************************
//FUNCTION:
variable_list CProjectGaussians::forward(AutogradContext* vCtx,
    torch::Tensor vMeans,
    torch::Tensor vScales,
    float vGlobScale,
    torch::Tensor vQuats,
    torch::Tensor vViewMat,
    torch::Tensor vProjMat,
    float vFx,
    float vFy,
    float vCx,
    float vCy,
    int vImgHeight,
    int vImgWidth,
    TileBounds vTileBounds,
    float vClipThresh
) {

    int NumPoints = vMeans.size(0);

    auto ForwardTensor = project_gaussians_forward_tensor(NumPoints, vMeans, vScales, vGlobScale,
        vQuats, vViewMat, vProjMat, vFx, vFy,
        vCx, vCy, vImgHeight, vImgWidth, vTileBounds, vClipThresh);
    torch::Tensor Cov3d = std::get<0>(ForwardTensor);
    torch::Tensor XYs = std::get<1>(ForwardTensor);
    torch::Tensor Depths = std::get<2>(ForwardTensor);
    torch::Tensor Radii = std::get<3>(ForwardTensor);
    torch::Tensor Conics = std::get<4>(ForwardTensor);
    torch::Tensor NumTilesHit = std::get<5>(ForwardTensor);

    vCtx->saved_data["imgHeight"] = vImgHeight;
    vCtx->saved_data["imgWidth"] = vImgWidth;
    vCtx->saved_data["numPoints"] = NumPoints;
    vCtx->saved_data["globScale"] = vGlobScale;
    vCtx->saved_data["fx"] = vFx;
    vCtx->saved_data["fy"] = vFy;
    vCtx->saved_data["cx"] = vCx;
    vCtx->saved_data["cy"] = vCy;
    vCtx->save_for_backward({ vMeans, vScales, vQuats, vViewMat, vProjMat, Cov3d, Radii, Conics });

    return { XYs, Depths, Radii, Conics, NumTilesHit, Cov3d };
}

//******************************************************************
//FUNCTION:
tensor_list CProjectGaussians::backward(AutogradContext* vCtx, tensor_list vGradOutputs) {
    torch::Tensor InputXYs = vGradOutputs[0];
    torch::Tensor InputDepths = vGradOutputs[1];
    torch::Tensor InputRadii = vGradOutputs[2];
    torch::Tensor InputConics = vGradOutputs[3];
    torch::Tensor InputNumTiles = vGradOutputs[4];
    torch::Tensor InputCov3d = vGradOutputs[5];

    variable_list SavedList = vCtx->get_saved_variables();
    torch::Tensor Means = SavedList[0];
    torch::Tensor Scales = SavedList[1];
    torch::Tensor Quats = SavedList[2];
    torch::Tensor ViewMat = SavedList[3];
    torch::Tensor ProjMat = SavedList[4];
    torch::Tensor Cov3d = SavedList[5];
    torch::Tensor Radii = SavedList[6];
    torch::Tensor Conics = SavedList[7];

    auto BackwardTensor = project_gaussians_backward_tensor(vCtx->saved_data["numPoints"].toInt(),
        Means, Scales, vCtx->saved_data["globScale"].toDouble(),
        Quats, ViewMat, ProjMat,
        vCtx->saved_data["fx"].toDouble(), vCtx->saved_data["fy"].toDouble(),
        vCtx->saved_data["cx"].toDouble(), vCtx->saved_data["cy"].toDouble(),
        vCtx->saved_data["imgHeight"].toInt(), vCtx->saved_data["imgWidth"].toInt(),
        Cov3d, Radii,
        Conics, InputXYs, InputDepths, InputConics);
    torch::Tensor None;

    return { std::get<2>(BackwardTensor), //Authorm comment:  v_mean
            std::get<3>(BackwardTensor), //Authorm comment:  v_scale
            None, //Authorm comment:  globScale
            std::get<4>(BackwardTensor), //Authorm comment:  v_quat
            None, //Authorm comment:  viewMat
            None, //Authorm comment:  projMat
            None, //Authorm comment:  fx
            None, //Authorm comment:  fy
            None, //Authorm comment:  cx
            None, //Authorm comment:  cy
            None, //Authorm comment:  imgHeight
            None, //Authorm comment:  imgWidth
            None, //Authorm comment:  tileBounds
            None // Authorm comment: clipThresh
    };
}